import os.path
import time

import numpy as np
import cv2
import imageConversion
import matplotlib.pyplot as plt
from PIL import Image


def clean_height(image, nan_num, npy_name):
    """
    Auth: WZW
    根据nan，提取周围部分，并将剩余的nan数据转为nan_num
    UpArea: [ :1500, : ]
    SideArea: [ :2900, :1000]
    PreVersion: clean_height in V1
    Args:
        image: np.array
        nan_num:
    Returns:
        compressed data
    """
    if npy_name.count("_side_") > 0:
        image = image[:2900, :1000]
    else:
        image = image[:1700, :]
    left = 0
    right = 0
    row, column = image.shape
    for i in range(row):
        if len(image[i, np.isnan(image[i, :])]) != column:
            left = i
            break
    for i in range(row - 1, -1, -1):
        if len(image[i, np.isnan(image[i, :])]) != column:
            right = i + 1
            break
    image = image[left:right, :]
    row, column = image.shape
    left = 0
    right = 0
    for i in range(column):
        if len(image[np.isnan(image[:, i]), i]) != row:
            left = i
            break
    for i in range(column - 1, -1, -1):
        if len(image[np.isnan(image[:, i]), i]) != row:
            right = i + 1
            break
    image = image[:, left:right]
    print(image.shape)
    if not np.isnan(nan_num):
        print("数据压缩，修改nan为：" + str(nan_num))
        nan_index = np.isnan(image)
        image[nan_index] = nan_num
    return image


def clean_light(img_name, img_show=False):
    """
    亮度图清洗，去除黑色部分，只保留工件
    必须传入包括文件名的文件路径
    Args:
        img_name: image name with file path
        imshow:
        imsave:
        target_folder: folder to save image

    Returns:
        rgb image
    """
    image = cv2.imread(img_name)
    if img_name.count("_side_") > 0:
        image = image[:2900, :1000, :]
    else:
        image = image[:1700, :]
    """
    clean pic
    """
    img_gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    row, column = img_gray.shape
    left = 0
    right = 0
    for i in range(row):
        if np.sum(img_gray[i, :]) != 0:
            left = i
            break
    for i in range(row - 1, -1, -1):
        if np.sum(img_gray[i, :]) != 0:
            right = i + 1
            break
    img_gray = img_gray[left:right, :]
    image = image[left:right, :, :]
    row, column = img_gray.shape
    for i in range(column):
        if np.sum(img_gray[:, i]) != 0:
            left = i
            break
    for i in range(column - 1, -1, -1):
        if np.sum(img_gray[:, i]) != 0:
            right = i + 1
            break
    image = image[:, left:right, :]
    print(image.shape)
    if img_show:
        # cv2.imshow('img_cut_c', c)
        # cv2.imshow('rgb_c', img_c)
        # cv2.waitKey(0)
        plt.imshow(image)
        plt.show()
    return image


def image_narrow_8bit(img_light, img_height):
    """

    Args:
        img_light: nd.array, light image
        img_height: nd.array, height image
    Returns:

    """
    print(img_light.shape)
    print(img_height.shape)
    a, b = img_light.shape[:2]
    img_light = cv2.resize(img_height, (b, a))
    return img_light


def image_blend_alpha_8bit(img_rgb, img_height, alpha=0.5, img_show=False, img_save=False, target_folder="",
                           file_name="out"):
    """
    Auth: WZW
    Args:
        img_rgb: nd.array, light image
        img_height: nd.array, height image
        alpha: 0 - 1, used to blend picture

    Returns:

    """
    print(img_rgb.shape)
    print(img_height.shape)

    img_rgb_pil = Image.fromarray(img_rgb, mode='RGB')
    img_height_pil = Image.fromarray(img_height, mode='RGB')

    img_rgb_alpha = img_rgb_pil.convert('RGBA')
    img_height_alpha = img_height_pil.convert('RGBA')

    blend_img = Image.blend(img_height_alpha, img_rgb_alpha, alpha)
    # img_height_alpha.show()
    if img_show:
        img_rgb_alpha.show()
        blend_img.show()
    if img_save:
        img_rgb_alpha.save(os.path.join(target_folder, file_name + "_light.png"))
        blend_img.save(os.path.join(target_folder, file_name + "_fusion.png"))
    return blend_img


def image_channels_add_8bit(img_rgb, img_height, img_show=False, img_save=False, target_folder="./",
                            file_name="out"):
    """

    Args:
        img_rgb:
        img_height:
        img_show:
        img_save:
        target_folder:
        file_name:

    Returns:

    """
    meta_img = np.zeros(img_rgb.shape)
    img_height_gray = cv2.cvtColor(img_height, cv2.COLOR_RGB2GRAY)
    img_rgb_gray = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2GRAY)
    meta_img[:, :, 0] = img_height_gray
    meta_img[:, :, 1] = img_rgb_gray
    # meta_img[:, :, 2] = meta_img[:, :, 2] * 50
    meta_img = meta_img.astype(np.uint8)
    if img_show:
        # cv2.imshow("test", meta_img)
        # cv2.waitKey(0)
        plt.imshow(meta_img)
        plt.show()
    if img_save:
        meta_img = cv2.cvtColor(meta_img, cv2.COLOR_BGR2RGB)
        cv2.imwrite(os.path.join(target_folder, file_name + "_channelsBlend.png"), meta_img)
    return


def generate_dataset_image(filePath, imgPath, img_show=False, img_save=False, targetFolder="./"):
    """
    Auth: WZW
    Args:
        filePath: npy 文件路径，包括文件名及后缀
        imgPath: 图片文件路径，包括文件名及后缀
        img_show:
        img_save:
        targetFolder: 目标目录

    Returns:

    """
    img_light = clean_light(imgPath, img_show=img_show)
    npData = np.load(filePath)
    img_height = clean_height(npData, np.nan, filePath)
    img16, img8, img_max, img_min = imageConversion.npy_normalization_16_8(img_height, 0)
    img8_rgb = imageConversion.gray2rgb_8bit(img8)
    reshaped_img_height = image_narrow_8bit(img_light, img8_rgb)
    image_blend_alpha_8bit(img_light, reshaped_img_height, img_show=img_show, img_save=img_save,
                           target_folder=targetFolder, file_name=os.path.basename(filePath).split('.')[0])
    image_channels_add_8bit(img_light, reshaped_img_height, img_show=img_show, img_save=img_save,
                            target_folder=targetFolder, file_name=os.path.basename(filePath).split('.')[0])


if __name__ == "__main__":
    """
    测试高度图清洗    
    """
    side_npy = "D:\\Programs\\data\\dataset\\20220628-full-size\\2022-06-23\\Side\\35051772623505277134A23062200150556_side_2022-06-23-15-39-35.npy"
    side_png = "D:\\Programs\\data\\dataset\\20220628-full-size\\2022-06-23\\Side\\35051772623505277134A23062200150556_side_2022-06-23-15-39-35_meta.png"
    up_npy = "D:\\Programs\\data\\dataset\\20220628-full-size\\2022-06-23\\Up\\35051772623505277134A23062200150556_up_2022-06-23-15-39-06.npy"
    up_png = "D:\\Programs\\data\\dataset\\20220628-full-size\\2022-06-23\\Up\\35051772623505277134A23062200150556_up_2022-06-23-15-39-06_meta.png"
    # da = np.load(up_npy)
    # start = time.time()
    # img_height = clean_height(da, np.nan, up_png)
    # end = time.time()
    # print("+==========================+ " + str(end - start) + " +==========================+")
    # img16, img8, img_max, img_min = imageConversion.npy_normalization_16_8(img_height, 0, imgShow=True)
    # clean_light(img_name=up_png, img_show=True)

    """
    生成目标文件测试
    """
    generate_dataset_image(side_npy, side_png, img_save=True)
